# -*- coding:utf-8 -*-
# @Time: 2020/9/9 16:57
# @Author: gupeng
# @File: graph.py
# @Desc: 有向图--hexists,hmset,hmget,hkeys,hgetall测试
def make_edge_name_from_vertexs(start, end):
    """
    使用边的起点和终点组建边的名字
    :param start: 起点
    :param end: 终点
    :return: 边的名字
    """
    return str(start) + '->' + str(end)


def decompose_vertexs_from_edge_name(name):
    """
    从边的名字分解出起点和终点
    :param name: 边的名字'a->b'
    :return: ['a','b']
    """
    return name.split('->')


class Graph:

    def __init__(self, client, key):
        self.client = client
        self.key = key

    def add_edge(self, start, end, weight):
        """
        新增一条从start到end的边，权重为weight
        :param start: 起点
        :param end: 终点
        :param weight: 权重
        :return: None
        """
        edge = make_edge_name_from_vertexs(start, end)
        self.client.hset(self.key, edge, weight)

    def remove_edge(self, start, end):
        """
        删除从start到end的边
        :param start: 起点
        :param end: 终点
        :return: 删除成功返回True，否则返回False
        """
        edge = make_edge_name_from_vertexs(start, end)
        return self.client.hdel(self.key, edge)

    def get_edge_weight(self, start, end):
        """
        获取从start到end的边的权重
        :param start: 起点
        :param end: 终点
        :return: 权重
        """
        edge = make_edge_name_from_vertexs(start, end)
        return self.client.hget(self.key, edge)

    def has_edge(self, start, end):
        """
        判断从start到end的边是否存在
        :param start: 起点
        :param end: 终点
        :return: True存在，False不存在
        """
        edge = make_edge_name_from_vertexs(start, end)
        return self.client.hexists(self.key, edge)

    def add_multi_edges(self, *tuples):
        """
        一次向图中添加多条边，多个格式为(start,end,weight)
        :param tuples: 边信息
        :return: None
        """
        nodes_and_weights = {}
        for start, end, weight in tuples:
            edge = make_edge_name_from_vertexs(start, end)
            nodes_and_weights[edge] = weight
        self.client.hmset(self.key, nodes_and_weights)

    def get_multi_edge_weights(self, *tuples):
        """
        获取多个边的权重，多个格式为(start,end)
        :param tuples: 边信息
        :return: 多个边的权重
        """
        edge_list = []
        for start, end in tuples:
            edge = make_edge_name_from_vertexs(start, end)
            edge_list.append(edge)
        return self.client.hmget(self.key, edge_list)

    def get_all_edge(self):
        """
        获取所有的边
        :return: 所有的边
        """
        edges = self.client.hkeys(self.key)
        result = set()
        for edge in edges:
            start, end = decompose_vertexs_from_edge_name(edge.decode('utf-8'))
            result.add((start, end))
        return result

    def get_all_edge_with_weight(self):
        """
        返回所有的边以及他们的权重
        :return: 所有的边以及权重信息
        """
        edges_and_weights = self.client.hgetall(self.key)
        result = set()
        for edge, weight in edges_and_weights.items():
            start, end = decompose_vertexs_from_edge_name(edge.decode('utf-8'))
            result.add((start, end, weight))
        return result
